#include "lv_draw_img_png8_decoder.h"

#if USE_LV_DRAW_IMG_PNG8_DECODER

#include <stdio.h>
#include <stdbool.h>
#include "lvgl/lvgl.h"
#include "lvgl/src/lv_misc/lv_mem.h"
#include "lvgl/src/lv_draw/lv_draw_img.h"
#include "png.h"
#include "pngstruct.h"
#include "pnginfo.h"
#include "lv_drivers/display/scaler.h"

#define DEBUG_MEM 0

#define LV_PNG8_SIGNATURE_LEN 8
#define LV_PNG8_COLOR_TYPE_LEN 26

typedef struct {
    lv_fs_file_t * f;//at the end of open(), it will be closed, set NULL.
    int read_pos;
    png_structp png_ptr;
    png_infop info_ptr;
    png_bytep row_buf;
    int pixel_size_png;
    int pixel_size_lv;
    int read_line_mode;
    int row_curr;
} lv_png8_decoder_data_t;

#if DEBUG_MEM
static int malloc_size = 0;
#endif

/**********************
 *   STATIC FUNCTIONS
 **********************/
static png_voidp _png8_malloc(png_structp ptr, png_alloc_size_t size)
{
    png_voidp p;
    (void)ptr;

    p = lv_mem_alloc(size);

#if DEBUG_MEM
    int size_1 = size;
#ifdef LV_MEM_ENV64
    /*Round the size up to 8*/
    if(size_1 & 0x7) {
        size_1 = size_1 & (~0x7);
        size_1 += 8;
    }
#else
    /*Round the size up to 4*/
    if(size_1 & 0x3) {
        size_1 = size_1 & (~0x3);
        size_1 += 4;
    }
#endif
    if(p) {
        malloc_size += size_1;
        printf("png8_malloc: total: %d\n", malloc_size);
    }
#endif

    return p;
}

static void _png8_free(png_structp ptr, png_voidp p)
{
    (void)ptr;

#if DEBUG_MEM
    malloc_size -= lv_mem_get_size(p);
    printf("png8_free: total: %d\n", malloc_size);
#endif

    lv_mem_free(p);
    return;
}

static void png8_read_data_fn(png_structp ptr, png_bytep data,
                              png_size_t length)
{
    const uint8_t * p;
    lv_img_decoder_dsc_t * dsc = (lv_img_decoder_dsc_t *)ptr->io_ptr ;
    if(!dsc) return;
    lv_png8_decoder_data_t * user_data = dsc->user_data;
    if(!user_data) return;

    //file
    if(user_data->f) {
        /*file form romfs or psram*/
        uint32_t readbytes = 0;
        lv_fs_res_t res;

        res = lv_fs_seek(user_data->f, user_data->read_pos);
        if(res != LV_FS_RES_OK) {
            printf("png8_read_data_fn png decoder lv_fs_seek failed\n");
            return;
        }

        res = lv_fs_read(user_data->f, data, length, &readbytes);
        if(res != LV_FS_RES_OK || length != readbytes) {
            printf("png8_read_data_fn png decoder lv_fs_read failed\n");
        }
    } else if(dsc->src_type == LV_IMG_SRC_VARIABLE) {
        lv_img_dsc_t * img = (lv_img_dsc_t *)dsc->src;
        /* img->data is png data */
        p = img->data + user_data->read_pos;
        memcpy(data, p, length);
    }

    user_data->read_pos += length;
}

static void
png8_do_expand_index(png_structrp png_ptr, png_bytep row)
{
    int shift, value;
    png_bytep sp, dp;
    png_uint_32 i;
    png_uint_32 row_width = png_ptr->width;

    if(png_ptr->bit_depth < 8) {
        switch(png_ptr->bit_depth) {
            case 1: {
                sp = row + (size_t)((row_width - 1) >> 3);
                dp = row + (size_t)row_width - 1;
                shift = 7 - (int)((row_width + 7) & 0x07);
                for(i = 0; i < row_width; i++) {
                    if((*sp >> shift) & 0x01)
                        *dp = 1;

                    else
                        *dp = 0;

                    if(shift == 7) {
                        shift = 0;
                        sp--;
                    }

                    else
                        shift++;

                    dp--;
                }
                break;
            }

            case 2: {
                sp = row + (size_t)((row_width - 1) >> 2);
                dp = row + (size_t)row_width - 1;
                shift = (int)((3 - ((row_width + 3) & 0x03)) << 1);
                for(i = 0; i < row_width; i++) {
                    value = (*sp >> shift) & 0x03;
                    *dp = (png_byte)value;
                    if(shift == 6) {
                        shift = 0;
                        sp--;
                    }

                    else
                        shift += 2;

                    dp--;
                }
                break;
            }

            case 4: {
                sp = row + (size_t)((row_width - 1) >> 1);
                dp = row + (size_t)row_width - 1;
                shift = (int)((row_width & 0x01) << 2);
                for(i = 0; i < row_width; i++) {
                    value = (*sp >> shift) & 0x0f;
                    *dp = (png_byte)value;
                    if(shift == 4) {
                        shift = 0;
                        sp--;
                    }

                    else
                        shift += 4;

                    dp--;
                }
                break;
            }

            default:
                break;
        }
    }
}


static void png8_do_l8_image_data(unsigned char * pbuf, uint16_t num_palette, png_colorp palette, uint16_t num_trans, png_bytep trans_alpha)
{
    int i;
    unsigned char * p = pbuf;

    *p++ = 0x00;    // ARGB format
    *p++ = 0x00;    // reserved
    *p++ = (unsigned char)num_palette;    // low 8 bits of palette cnt
    *p++ = (unsigned char)(num_palette >> 8);   // high 8 bits of palette cnt
    for(i = 0; i < num_trans; i++) {
        *p++ = palette->red;
        *p++ = palette->green;
        *p++ = palette->blue;
        *p++ = *trans_alpha;
        palette++;
        trans_alpha++;
    }
    for(i = 0; i < num_palette - num_trans; i++) {
        *p++ = palette->red;
        *p++ = palette->green;
        *p++ = palette->blue;
        *p++ = 0xFF;
        palette++;
    }
}

/**
 * Open a PNG image and return the decided image
 * @param src can be file name or pointer to a C array
 * @param style style of the image object (unused now but certain formats might use it)
 * @return pointer to the decoded image or  `LV_RES_INV` if failed
 */
static lv_res_t png8_decoder_open(lv_img_decoder_t * decoder, lv_img_decoder_dsc_t * dsc)
{
    /*Open the file if it's a file*/
    png_bytepp row_pointers;
    unsigned width, height;
    png_bytep img_buf = NULL;


    lv_fs_file_t f;

    /*allocat user_data*/
    if(!dsc->user_data) {
        dsc->user_data = lv_mem_alloc(sizeof(lv_png8_decoder_data_t));
        if(!dsc->user_data) {
            printf("[%s]: out of memory", __func__);
            goto err;
        }
        memset(dsc->user_data, 0, sizeof(lv_png8_decoder_data_t));
    }
    lv_png8_decoder_data_t * user_data = dsc->user_data;

    /*If it's a PNG file...*/
    /*A PNG file, not support Zoom in, show original size*/
    /*
     *eg:
     *"C:/xxx.png", means file in PSRAM;
     *"R:/xxx.png", means file in ROMFS XIP NOR Flash;
     *"P:/xxx.png", means file in PC;
     */
    if(dsc->src_type == LV_IMG_SRC_FILE) {
        const char * fn = (char *)dsc->src;

        /*Support only "*.png8" files*/
        if(strcmp(lv_fs_get_ext(fn), "png8")) {
            printf("[%s]: file %s is not .png8\n", __func__, fn);
            goto err;
        }

        lv_fs_res_t ret = lv_fs_open(&f, fn, LV_FS_MODE_RD);
        if(ret != LV_FS_RES_OK) {
            printf("[%s]: lv_fs_open failed\n", __func__);
            goto err;
        }
        user_data->f = lv_mem_alloc(sizeof(f));
        if(user_data->f == NULL) {
            printf("[%s]: out of memory", __func__);
            goto err;
        }
        memcpy(user_data->f, &f, sizeof(f));
    } else {
        printf("[%s]: not support src_type\n", __func__);
        goto err;
    }

    user_data->read_pos = LV_PNG8_SIGNATURE_LEN;//skip sig
    user_data->png_ptr = png_create_read_struct_2(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL, NULL, _png8_malloc, _png8_free);
    if(!user_data->png_ptr) {
        printf("png_create_read_struct_2 failed\n");
        goto err;
    }

    user_data->info_ptr = png_create_info_struct(user_data->png_ptr);
    if(!user_data->info_ptr) {
        printf("png_create_info_struct failed\n");
        png_destroy_read_struct(&user_data->png_ptr, (png_infopp)NULL, (png_infopp)NULL);
        goto err;
    }

    if(setjmp(png_jmpbuf(user_data->png_ptr))) {
        printf("set jmp failed\n");
        png_destroy_read_struct(&user_data->png_ptr, (png_infopp)NULL, (png_infopp)NULL);
        goto err;
    }

    png_set_read_fn(user_data->png_ptr, (png_voidp)dsc, png8_read_data_fn);
    png_set_sig_bytes(user_data->png_ptr, LV_PNG8_SIGNATURE_LEN);
    png_read_info(user_data->png_ptr, user_data->info_ptr);


    /* png8 decoder does not need to expend */
//    if(user_data->info_ptr->color_type == PNG_COLOR_TYPE_PALETTE) {  /* PNG8 */
//        png_set_palette_to_rgb(user_data->png_ptr);
//    }

    if(!(user_data->info_ptr->trans_alpha))
    {
        png_set_strip_alpha(user_data->png_ptr);
    }

    dsc->header.cf = LV_IMG_CF_L8_ARGB8888;

    width = user_data->info_ptr->width;
    height = user_data->info_ptr->height;

    row_pointers = _png8_malloc(user_data->png_ptr, height * sizeof(png_bytep));
    //img_buf = _png_malloc(user_data->png_ptr, width * height * user_data->pixel_size_png);
    img_buf = _png8_malloc(user_data->png_ptr, width * height + 4 + (user_data->png_ptr->num_palette * 4));   //pixel index + 4 bytes format + ARGB palette

    //only no memeory, enter readline mode. normally is whole decoded mode which support cached.
    if((!row_pointers) || (!img_buf)) {
        //printf("no mem for open, try to readline\n");
        _png8_free(user_data->png_ptr, row_pointers);
        _png8_free(user_data->png_ptr, img_buf);
        img_buf = NULL;

        if((dsc->header.w != width) || (dsc->header.h != height)) {
            printf("no scale for readline, png8 decoder failed\n");
            png_destroy_read_struct(&user_data->png_ptr, &user_data->info_ptr, NULL);
            goto err;
        }

        //user_data->row_buf = _png8_malloc(user_data->png_ptr, (width * user_data->pixel_size_png));
        user_data->row_buf = _png8_malloc(user_data->png_ptr, width);
        if(!user_data->row_buf) {
            printf("no mem for readline, png8 decoder failed\n");
            png_destroy_read_struct(&user_data->png_ptr, &user_data->info_ptr, NULL);
            goto err;
        }
        user_data->read_line_mode = 1;
        user_data->row_curr = 0;
        return LV_RES_OK;
    }

    //whole decoded to bitmap, saved to cache.
    for(int i = 0; i < height; i++) {
        //row_pointers[i] = img_buf + i * width * user_data->pixel_size_png;
        row_pointers[i] = img_buf + i * width;
    }

    png_set_rows(user_data->png_ptr, user_data->info_ptr, row_pointers);
    png_read_image(user_data->png_ptr, row_pointers);
    png_read_end(user_data->png_ptr, user_data->info_ptr);

    for(int i = 0; i < height; i++) {
        png8_do_expand_index(user_data->png_ptr, row_pointers[i]);
    }
    png8_do_l8_image_data(img_buf + (width * height), user_data->png_ptr->num_palette, user_data->png_ptr->palette,
                          user_data->png_ptr->num_trans, user_data->png_ptr->trans_alpha
                         );


    _png8_free(user_data->png_ptr, row_pointers);
    png_destroy_read_struct(&user_data->png_ptr, &user_data->info_ptr, NULL);


    //total bit map saved in cache;
    dsc->img_data = img_buf;

    //close fd
    if(user_data->f) {
        lv_fs_close(user_data->f);
        lv_mem_free(user_data->f);
        user_data->f = NULL;
    }

    return LV_RES_OK;

err:
    if(dsc->user_data) {
        if(user_data->f) {
            lv_fs_close(user_data->f);
            lv_mem_free(user_data->f);
        }
        lv_mem_free(user_data);
        dsc->user_data = NULL;
    }
    return LV_RES_INV;
}

static void png8_decoder_close(lv_img_decoder_t * decoder, lv_img_decoder_dsc_t * dsc)
{
    (void)decoder; /*Unused*/

    lv_png8_decoder_data_t * user_data = (lv_png8_decoder_data_t *)dsc->user_data;
    if(user_data) {
        if(user_data->f) {
            lv_fs_close(user_data->f);
            lv_mem_free(user_data->f);
        }

        if(user_data->read_line_mode) {
            user_data->read_line_mode = 0;
            png_read_end(user_data->png_ptr, user_data->info_ptr);
            _png8_free(user_data->png_ptr, (png_voidp)user_data->row_buf);
            user_data->png_ptr = NULL;
            png_destroy_read_struct(&user_data->png_ptr, &user_data->info_ptr, NULL);
        } else if(dsc->img_data != NULL) {
            _png8_free(NULL, (png_voidp)dsc->img_data);
            dsc->img_data = NULL;
        }

        lv_mem_free(user_data);
        dsc->user_data = NULL;
    }
}

/**
 * Decode `len` pixels starting from the given `x`, `y` coordinates and store them in `buf`.
 * Required only if the "open" function can't return with the whole decoded pixel array.
 * @param decoder pointer to the decoder the function associated with
 * @param dsc pointer to decoder descriptor
 * @param x start x coordinate
 * @param y start y coordinate
 * @param len number of pixels to decode
 * @param buf a buffer to store the decoded pixels
 * @return LV_RES_OK: ok; LV_RES_INV: failed
 */
static lv_res_t png8_decoder_read_line(lv_img_decoder_t * decoder, lv_img_decoder_dsc_t * dsc, lv_coord_t x,
                                       lv_coord_t y, lv_coord_t len, uint8_t * buf)
{
    (void)decoder; /*Unused*/

    unsigned width;
    int i;
    width = dsc->header.w;

    lv_png8_decoder_data_t * user_data = dsc->user_data;
    if(!user_data) {
        printf("png8_decoder_read_line failed for empty user_data\n");
        return LV_RES_INV;
    }

    if(user_data->info_ptr->interlace_type != PNG_INTERLACE_NONE) {
        printf("error: readline does not support interlaced png8\n");
        return LV_RES_INV;
    }

    if(!user_data->row_buf) {
        printf("error: no mem for readline\n");
        return LV_RES_INV;
    }

    for(i = user_data->row_curr; i <= y; i++) {
        png_read_row(user_data->png_ptr, user_data->row_buf, NULL);
        user_data->row_curr++;
    }
    /*
        if(dsc->header.cf == LV_IMG_CF_RAW_ALPHA) {
            lv_color_rgba8888_to_16_alpha((lv_color_rgba8888_t *)user_data->row_buf, (lv_color16_alpha_t *)user_data->row_buf, width);
        } else {
            lv_color_rgb888_to_16((lv_color_rgb888_t *)user_data->row_buf, (lv_color16_t *)user_data->row_buf, width);
        }

        memcpy(buf, user_data->row_buf + x, len * user_data->pixel_size_lv);
    */
    png8_do_expand_index(user_data->png_ptr, user_data->row_buf);
    memcpy(buf, user_data->row_buf + x, len);
    png8_do_l8_image_data(buf + len, user_data->png_ptr->num_palette, user_data->png_ptr->palette,
                          user_data->png_ptr->num_trans, user_data->png_ptr->trans_alpha
                         );

    return LV_RES_OK;
}

/**
 * Get info about a PNG8 image
 * @param src can be file name or pointer to a C array
 * @param header store the info here
 * @return LV_RES_OK: no error; LV_RES_INV: can't get the info
 */
static lv_res_t png8_decoder_info(struct _lv_img_decoder * decoder, const void * src, lv_img_header_t * header)
{
    (void) decoder; /*Unused*/
    lv_img_src_t src_type = lv_img_src_get_type(src);          /*Get the source type*/

    /*If it's a PNG file...*/
    /*A PNG file, not support Zoom in, show original size*/
    /*
     *eg:
     *"C:/xxx.png", means file in PSRAM;
     *"R:/xxx.png", means file in ROMFS XIP NOR Flash;
     *"C:/xxx.png", means file in PC;
     */
    if(src_type == LV_IMG_SRC_FILE) {
        const char * fn = src;
        /*Support only "*.png8" files*/
        if(strcmp(lv_fs_get_ext(fn), "png8")) return LV_RES_INV;

        lv_fs_res_t ret;
        lv_fs_file_t f;
        ret = lv_fs_open(&f, fn, LV_FS_MODE_RD);
        if(ret != LV_FS_RES_OK) {
            printf("[%s]: lv_fs_open failed\n", __func__);
            return LV_RES_INV;
        }

        ret = lv_fs_seek(&f, 16);
        if(ret != LV_FS_RES_OK) {
            printf("[%s]: lv_fs_open failed\n", __func__);
            return LV_RES_INV;
        }

        uint32_t size[2];
        uint32_t br;
        ret = lv_fs_read(&f, (void *)size, 8, &br);
        if(ret != LV_FS_RES_OK || br != 8) {
            printf("[%s]: lv_fs_read failed\n", __func__);
            return LV_RES_INV;
        }

        uint8_t info[2];
        ret = lv_fs_read(&f, (void *)info, 2, &br);
        if(ret != LV_FS_RES_OK || br != 2) {
            printf("[%s]: lv_fs_read failed\n", __func__);
            return LV_RES_INV;
        }
        if(3 != info[1]) {
            printf("[%s]: failed non png8 file\n", __func__);
            lv_fs_close(&f);
            return LV_RES_INV;
        }

        // printf("%s bit depth is %d\n", fn,info[0]);

        lv_fs_close(&f);

        /*Save the data in the header*/
        header->always_zero = 0;
        header->cf = LV_IMG_CF_RAW_ALPHA;
        header->reserved = 1; //if reserved = 1. means file
        /*The width and height are stored in Big endian format so convert them to little endian*/
        header->w = (lv_coord_t)((size[0] & 0xff000000) >> 24) + ((size[0] & 0x00ff0000) >> 8);
        header->h = (lv_coord_t)((size[1] & 0xff000000) >> 24) + ((size[1] & 0x00ff0000) >> 8);

        return LV_RES_OK;

    }
    /* do not handle this type */
    else if(src_type == LV_IMG_SRC_VARIABLE) {
        return LV_RES_INV;
    }

    return LV_RES_INV;         /*If didn't succeeded earlier then it's an error*/
}

/**
 * Register the PNG decoder functions in LittlevGL
 */
void lv_png8_decoder_init(void)
{
    lv_img_decoder_t * dec = lv_img_decoder_create();
    lv_img_decoder_set_info_cb(dec, png8_decoder_info);
    lv_img_decoder_set_open_cb(dec, png8_decoder_open);
    lv_img_decoder_set_read_line_cb(dec, png8_decoder_read_line);
    lv_img_decoder_set_close_cb(dec, png8_decoder_close);
}

#endif /* USE_LV_DRAW_IMG_PNG8_DECODER */
